// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include "grpc++/channel.h" // @manual
#include "grpc++/client_context.h" // @manual
#include "grpc++/create_channel.h" // @manual
#include "grpc++/security/credentials.h" // @manual
#include "spark/connect/base.grpc.pb.h"
#include <pollux/common/utils.h>
#include <pollux/common/memory/memory.h>
#include <pollux/plan/reference_query_runner.h>
#include <pollux/vector/complex_vector.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::functions::sparksql::fuzzer {

using kumo::pollux::fuzzer::DataSpec;

/// Query runner that uses Spark as a reference database. It converts Pollux
/// query plan to Spark SQL and executes it in Spark. The results are returned
/// as Pollux compatible format.
class SparkQueryRunner : public ReferenceQueryRunner {
 public:
  /// @param coordinatorUri Spark connect server endpoint, e.g. localhost:15002.
  SparkQueryRunner(
      memory::MemoryPool* pool,
      const std::string& coordinatorUri,
      const std::string& userId,
      const std::string& userName)
      : ReferenceQueryRunner(pool),
        userId_(userId),
        userName_(userName),
        sessionId_(generateUUID()),
        stub_(spark::connect::SparkConnectService::NewStub(grpc::CreateChannel(
            coordinatorUri,
            grpc::InsecureChannelCredentials()))) {
    pool_ = aggregatePool()->addLeafChild("leaf");
    copyPool_ = aggregatePool()->addLeafChild("copy");
  };

  /// Converts Pollux query plan to Spark SQL. Supports Values -> Aggregation.
  /// Values node is converted into reading from 'tmp' table.
  /// @return std::nullopt for unsupported cases.
  std::optional<std::string> toSql(
      const pollux::core::PlanNodePtr& plan) override;

  std::multiset<std::vector<pollux::variant>> execute(
      const std::string& sql,
      const std::vector<pollux::RowVectorPtr>& input,
      const pollux::RowTypePtr& resultType) override;

  std::multiset<std::vector<pollux::variant>> execute(
      const std::string& sql,
      const std::vector<RowVectorPtr>& probeInput,
      const std::vector<RowVectorPtr>& buildInput,
      const RowTypePtr& resultType) override {
    POLLUX_NYI("SparkQueryRunner does not support join node.");
  }

  RunnerType runnerType() const override {
    return RunnerType::kSparkQueryRunner;
  }

  const std::vector<TypePtr>& supportedScalarTypes() const override;

  const std::unordered_map<std::string, DataSpec>&
  aggregationFunctionDataSpecs() const override;

  bool supportsPolluxVectorResults() const override {
    return true;
  }

  std::vector<pollux::RowVectorPtr> executeVector(
      const std::string& sql,
      const std::vector<pollux::RowVectorPtr>& input,
      const pollux::RowTypePtr& resultType) override;

  std::vector<pollux::RowVectorPtr> execute(const std::string& sql) override;

 private:
  using ReferenceQueryRunner::toSql;

  // Generates a random UUID string for Spark. It must be of the format
  // '00112233-4455-6677-8899-aabbccddeeff'.
  std::string generateUUID();

  pollux::memory::MemoryPool* pool() {
    return pool_.get();
  }

  // Reads the arrow IPC-format string data with arrow IPC reader and convert
  // them into Pollux RowVectors.
  std::vector<pollux::RowVectorPtr> readArrowData(const std::string& data);

  google::protobuf::Arena arena_;
  const std::string userId_;
  const std::string userName_;
  const std::string sessionId_;
  // Used to make gRPC calls to the SparkConnectService.
  std::unique_ptr<spark::connect::SparkConnectService::Stub> stub_;
  std::shared_ptr<pollux::memory::MemoryPool> pool_;
  std::shared_ptr<pollux::memory::MemoryPool> copyPool_;
};
} // namespace kumo::pollux::functions::sparksql::fuzzer
